#ifndef _FEMSPACE_H_
#define _FEMSPACE_H_

#include "Dofs.h"
#include "Mesh.h"
#include "Point.h"
#include "TemplateElement.h"
#include "AMG.cpp"
#include "Equation.h"
#include <Eigen/Sparse>
#include <Eigen/IterativeLinearSolvers>
#include <Eigen/Dense>
#include <ctime>
#include <cmath>
#include <set>
/**
 * @file   FEMSpace.h
 * @author AOM
 * @date   Sat Mar  4 10:24:15 2023
 *
 * @brief  有限元空间基类
 *         主要功能：生成网格，分配自由度，组装刚度矩阵、右端项，处理边界条件，求解。
 *
 *
 */

typedef Eigen::SparseMatrix<double> SpMat;
typedef Eigen::VectorXd RHS;
typedef Eigen::VectorXd Solution;

template <unsigned int DIM>
std::map<int, double> interpolateBoundaryValues(auto &_boundary_function, std::set<Dofs<DIM>> _boundary_dofs, std::vector<Point<DIM>> _vertex);

void applyBoundaryValues(std::map<int, double> &boundary_values, Eigen::SparseMatrix<double> &system_matrix, Eigen::VectorXd &system_rhs);

template <unsigned int DIM>
class FEMSpace
{
public:
    // virtual void make_grid() = 0;
    virtual void setup_system() = 0;
    virtual void assemble_system() = 0;
    virtual void dealwithBoundaryCondition() = 0;
    virtual void solve() = 0;
};

class Possion_2D : public FEMSpace<2>
{
public:
    Possion_2D();
    Possion_2D(Mesh<2> *, TemplateElement<2> *, Equation<2> *, Domain<2> *);
    Possion_2D(const Possion_2D &);
    ~Possion_2D(){/*delete _mesh;delete _element;delete _eqation;*/};
    // virtual void make_grid();
    virtual void setup_system();
    virtual void assemble_system();
    virtual void dealwithBoundaryCondition();
    virtual void solve();

    double LocalStiffMatrix(int i, int j);
    Solution getSolution() const;
    SpMat A() const;
    RHS Rhs() const;

protected:
    SpMat system_matrix;
    RHS system_rhs;
    Solution solution;
    Domain<2> *domain;
    Mesh<2> *mesh;
    TemplateElement<2> *element;
    Equation<2> *eqation;
    // BoundaryCondition<2> _boundarycondition;
};

double Possion_2D::LocalStiffMatrix(int i, int j)
{
    double a = 0;
    for (int k = 0; k < element->n_GaussPnt(); k++)
    {
        double xi = element->GaussionPoint(k)[0];
        double eta = element->GaussionPoint(k)[1];
        double b = 0;
        for (int l = 0; l < element->gradient(xi, eta, i).size(); l++)
            b += element->gradient(xi, eta, i)[l] * element->gradient(xi, eta, j)[l];
        a = a + element->det_Jacobi(xi, eta) * element->GaussionWeight(k) * b;
    }
    return a;
}

Solution Possion_2D::getSolution() const
{
    return solution;
}

Possion_2D::Possion_2D()
{
    int num = mesh->getTotalNumDofs();
    system_matrix = SpMat(num, num);
    system_rhs = RHS(num);
    //_boundarycondition.apply(mesh);
}
Possion_2D::Possion_2D(Mesh<2> *_mesh, TemplateElement<2> *_element, Equation<2> *_equation, Domain<2> *_domain)
{
    domain = _domain;
    mesh = _mesh;
    element = _element;
    eqation = _equation;
    //_boundarycondition = (BoundaryCondition);
    int num = mesh->getTotalNumDofs();
    system_matrix = SpMat(num, num);
    system_rhs = RHS(num);
    //_boundarycondition.apply(mesh);
}

Possion_2D::Possion_2D(const Possion_2D &p)
{
    domain = p.domain;
    mesh = p.mesh;
    element = p.element;
    eqation = p.eqation;
    //_boundarycondition = p._boundarycondition;
    system_matrix = p.system_matrix;
    system_rhs = p.system_rhs;
    //_boundarycondition.apply(_mesh);
}

void Possion_2D::setup_system()
{
    /*Enumerating is done by using DistributeDofs(),2-order element*/
    mesh->distributeDofs(); /*<  Q1网格对应Q1单元，Q2网格对应Q2单元.关注成员变量是否对应 */
    mesh->setBoundaryIndex();
    std::cout << "Number of degrees of freedom: " << mesh->getTotalNumDofs() << std::endl; /*<  changed with order,same for below,same for getBoundaryDofs(2)  */
    //_element->setMesh(*mesh);

    /**
     * set the sizes of system_matrix, right hand side vector ans the solution vector;
     */
    int num = mesh->getTotalNumDofs();
    system_matrix = Eigen::SparseMatrix<double>(num, num);
    system_matrix.setZero();
    system_rhs = Eigen::VectorXd::Zero(num);
    solution = Eigen::VectorXd::Zero(num);
}

void Possion_2D::assemble_system()
{
    /**
     * Get the right_hand_side function and boundary function
     */
    auto right_hand_side = eqation->RightHandsTermFunction();
    // std::function<double(Dofs<2U>&)> boundary_function = BoundaryFunction;

    /**
     * for each Grid, we construct a cell_matirx and cell_rhs.
     */
    const unsigned int dofs_per_cell = element->n_Dofs();
    vector<Eigen::Triplet<double>> TriList(mesh->getGrid().size() * dofs_per_cell * dofs_per_cell);
    vector<Eigen::Triplet<double>>::iterator it = TriList.begin();
    Eigen::MatrixXd cell_matrix(dofs_per_cell, dofs_per_cell);
    Eigen::VectorXd cell_rhs(dofs_per_cell);

    /*assemble the global matrix and vector cell-by-cell.*/
    for (auto &cell : mesh->getGrid())
    {
        // bind the TemplateElement with current cell.

        element->reinit(cell, *mesh);

        cell_matrix = Eigen::MatrixXd::Zero(dofs_per_cell, dofs_per_cell);
        cell_rhs = Eigen::VectorXd::Zero(dofs_per_cell);

        for (int i = 0; i < dofs_per_cell; i++)
        {
            for (int j = 0; j < dofs_per_cell; j++)
            {
                *it = Eigen::Triplet<double>(element->getGlobalIndex(i), element->getGlobalIndex(j), LocalStiffMatrix(i, j));
                it++;
            }
        }

        //     Real a = 0;
        //     for (int k = 0; k < _element->n_GaussPnt(); k++)
        //     {
        // 	Real xi = _element->GaussionPoint(k)[0];
        // 	Real eta = _element->GaussionPoint(k)[1];
        //     a = a + _element->det_Jacobi(xi, eta) * _element->GaussionWeight(k)
        // * InnerProuduct(_element->gradient(xi,eta,i),_element->gradient(xi,eta,j));
        //
        //     }
        //     return static_cast<double>(a);
        for (int q_index = 0; q_index < element->n_GaussPnt(); ++q_index)
        {

            Point<2> pnt = element->GaussionPoint(q_index);

            // for (int i = 0; i < dofs_per_cell; ++i)
            // {
            // for (int j = 0; j < dofs_per_cell; ++j)
            // {
            //     double InnerProduct = 0;
            //     for(int k = 0;k < _element->gradient(pnt[0],pnt[1],i+1).size();k++)
            //     {

            // 	InnerProduct += _element->gradient(pnt[0],pnt[1],i+1)[k] *
            // 	    _element->gradient(pnt[0],pnt[1],j+1)[k];
            //     }
            //     cell_matrix(i, j) +=(InnerProduct *
            // 			 _element->det_Jacobi(pnt[0], pnt[1]) *
            // 			 _element->GaussionWeight(q_index));
            // }
            // }
            double xi = pnt[0], eta = pnt[1];
            double x = element->Global_x(xi, eta);
            double y = element->Global_y(xi, eta);
            double p[2] = {x, y};
            for (int i = 0; i < dofs_per_cell; ++i)
            {

                // Point<2> p({x,y});

                // Dofs<2> globalDof(p);
                cell_rhs(i) += (element->phi(xi, eta, i) *
                                right_hand_side(p) *
                                element->det_Jacobi(xi, eta) *
                                element->GaussionWeight(q_index));
            }
        }

        // for (int i = 0; i < dofs_per_cell; ++i)
        // {
        //     for (int j = 0; j < dofs_per_cell; ++j)
        //     {
        // 	int row = _element->GetGlobalIndex(i+1);
        // 	int column = _element->GetGlobalIndex(j+1);
        // 	system_matrix.coeffRef(row, column)+=cell_matrix(i, j);
        //     }
        // }
        for (int i = 0; i < dofs_per_cell; ++i)
            system_rhs[element->getGlobalIndex(i)] += cell_rhs(i);
    }
    system_matrix.setFromTriplets(TriList.begin(), TriList.end());
    system_matrix.makeCompressed();
    // cout << system_matrix << endl;
}

void Possion_2D::dealwithBoundaryCondition()
{
    std::set<Dofs<2>> boundary_dofs;
    for (auto &grid : mesh->getBoundaryGrid())
    {
        element->reinit(grid, *mesh);
        for (auto &dof : element->getBoundaryDofsonGrid())
            boundary_dofs.insert(dof);
    }
    // for (auto &it : boundary_dofs)
    //     cout << it << endl;
    // sort(boundary_dofs.begin(), boundary_dofs.end());
    // boundary_dofs.erase(unique(boundary_dofs.begin(), boundary_dofs.end()), boundary_dofs.end());
    std::map<int, double> boundary_values;
    auto boundary_function = eqation->BoundaryConditionFunction();
    std::vector<Point<2>> vertex = {domain->getVertex(0), domain->getVertex(2)};
    boundary_values = interpolateBoundaryValues(boundary_function, boundary_dofs, vertex);
    applyBoundaryValues(boundary_values, system_matrix, system_rhs);
}

void Possion_2D::solve()
{
    /*<  CG solve */
    Eigen::ConjugateGradient<Eigen::SparseMatrix<double>> Solver_sparse;
    Solver_sparse.setTolerance(1e-20);
    Solver_sparse.compute(system_matrix);
    solution = Solver_sparse.solve(system_rhs);
    /*<  AMG solve */
    // AMG_Multigrid AmgSolve(system_matrix, system_rhs, 1e-12, GS, 2.0 / 3, 0.5);
    // AmgSolve.Solve();
    // solution = AmgSolve.GetResult();
}

SpMat Possion_2D::A() const
{
    return system_matrix;
}
RHS Possion_2D::Rhs() const
{
    return system_rhs;
}

/**
 * @brief 返回map，存放边界上自由度的全局标号，和边界函数(仅Dirichlet)在对应自由度上的函数值。
 *
 * @param _boundary_function 边界函数
 * @param _boundary_dofs 边界上的自由度
 */
template <unsigned int DIM>
std::map<int, double> interpolateBoundaryValues(auto &_boundary_function, std::set<Dofs<DIM>> _boundary_dofs, std::vector<Point<DIM>> _vertex)
{
    std::map<int, double> boundary_values;
    for (auto &dofs : _boundary_dofs)
    {
        int index = dofs.getGlobalIndex();
        double p[DIM];
        for (int i = 0; i < DIM; i++)
            p[i] = dofs[i];
        double value;
        if (p[1] == _vertex[0][1])
            value = _boundary_function[0](p);
        else if (p[0] == _vertex[1][0])
            value = _boundary_function[1](p);
        else if (p[1] == _vertex[1][1])
            value = _boundary_function[2](p);
        else if (p[0] == _vertex[0][0])
            value = _boundary_function[3](p);
        else
            abort(3);
        boundary_values.insert({index, value});
    }
    return boundary_values;
}
/**
 * @brief 对刚度矩阵和右端项应用Dirichlet条件
 *
 * @param system_matrix 刚度矩阵
 * @param system_rhs 右端项
 */

void applyBoundaryValues(std::map<int, double> &boundary_values, Eigen::SparseMatrix<double> &system_matrix, Eigen::VectorXd &system_rhs)
{
    for (auto &it : boundary_values)
    {
        // 边界自由度的全局编号。
        int k = it.first;
        // 对应自由的边界函数值。
        double bnd_value = it.second;
        system_rhs[k] = bnd_value * system_matrix.coeffRef(k, k);
        for (Eigen::SparseMatrix<double>::InnerIterator it(system_matrix, k); it; ++it)
        {
            int row = it.row();
            if (row == k)
                continue;
            system_rhs[row] -= system_matrix.coeffRef(k, row) * bnd_value;
            system_matrix.coeffRef(k, row) = 0.0;
            system_matrix.coeffRef(row, k) = 0.0;
        }
    }
}

#else
#endif
